#pragma once

#include <string>
#include <unordered_map>
#include <vector>
#include <cstdint>
#include <unordered_set>
#include <memory>

namespace sca {
    class CFGBuilder {
        // define a type to describe an instruction
        using Instruction = std::string;
        // define a type to describe an instruction address
        using Address = std::string;
    public:

        class BasicBlock {
        private:
            std::vector<std::shared_ptr<Instruction>> instructions;
            std::vector<std::shared_ptr<BasicBlock>> nextbbs;
        public:
            BasicBlock() = default;

            ~BasicBlock() = default;

            void addInstructionIntoBB(std::shared_ptr<Instruction> &instruction) {
                instructions.push_back(instruction);
            }

            void addNextBasicBlock(std::shared_ptr<BasicBlock> &bb) {
                nextbbs.push_back(bb);
            }

            std::vector<std::shared_ptr<BasicBlock>> &getNextBasicBlocks() {
                return nextbbs;
            }

            std::shared_ptr<Instruction> getInstruction(size_t i) {
                if (i >= instructions.size()) {
                    return nullptr;
                }
                return instructions[i];
            }

            // 获取基本块的最后一条指令
            std::shared_ptr<Instruction> getTheLastInstruction() {
                if (instructions.empty()) {
                    return nullptr;
                }
                return getInstruction(instructions.size() - 1);
            }

            // 获取基本块的指令数量
            size_t getNumOfInstructions() {
                return instructions.size();
            }

            std::vector<std::shared_ptr<BasicBlock>> getNextBBs() {
                return nextbbs;
            }

            std::vector<std::shared_ptr<Instruction>> getInstructions() {
                return instructions;
            }
        };


        std::string targetPath;     // 被测的目标文件路径
        std::string tempFileDir;    // 临时文件目录，用于存储反汇编的文件

        std::unordered_map<Address, std::shared_ptr<Instruction>> jInstruction;
        std::unordered_map<Address, std::shared_ptr<Instruction>> callInstruction;
        std::unordered_map<Address, std::shared_ptr<Instruction>> retInstruction;

        std::unordered_map<Address,std::vector<std::shared_ptr<Instruction>>> beginAddr2Ret;
        std::unordered_map<std::shared_ptr<Instruction>,std::vector<std::shared_ptr<Instruction>>> retInst2CallInsts;

        // 基本块与基本块首地址之间的映射
        std::vector<std::shared_ptr<BasicBlock>> bbs;
        std::unordered_map<std::shared_ptr<BasicBlock>, Address> bb2addr;
        std::unordered_map<Address, std::shared_ptr<BasicBlock>> addr2bb;


        // 指令与指令地址之间的映射
        std::vector<std::shared_ptr<Instruction>> ins;
        std::unordered_map<std::shared_ptr<Instruction>, Address> inst2Add;
        std::unordered_map<Address, std::shared_ptr<Instruction>> add2Inst;


        std::unordered_set<Address> entryAdds; // 记录所有的入口地址

    public:
        explicit CFGBuilder(const std::string &_targetPath, const std::string &_tempFileDir = "/home/zengao/sca-binary-file/tmp/");

        // 从汇编文件中加载指令
        void loadInstruction();

        // 获取所有基本块的入口地址
        void loadEntryAdds();

        // 根据入口地址划分基本块
        bool buildBasicBlocksVec();

        // 添加所有基本块的依赖关系
        void createDependencies();

        void buildCallAndRetRelationship();

        std::vector<std::vector<int>> generationCFG();

        bool isTargetExist() const;
    };
}